from langchain.memory import ConversationBufferMemory
from langchain_community.chat_message_histories import RedisChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
#from langsmith.run_helpers import R
from .Prompt import PromptClass
from dotenv import load_dotenv

load_dotenv()
import os
import logging

logger = logging.getLogger(__name__)

redis_url = os.getenv("REDIS_URL", "redis://localhost:6380/0")
print(f"Redis URL: {redis_url}")


class MemoryClass:
    def __init__(self, memoryKey = "chat_history", model = os.getenv("BASE_MODEL")):
        self.memoryKey = memoryKey
        self.memory = []
        self.chatmodel = ChatOpenAI(model=model, base_url=os.getenv("OPENAI_API_BASE"))



    def summary_chain(self, store_message):
        try:
            prompt_instance = PromptClass()
            SystemPrompt = prompt_instance.system_prompt.format(feelScore=5, who_you_are="")
            Moods = prompt_instance.MOODS
            prompt = ChatPromptTemplate.from_messages(
                [
                    ("system", SystemPrompt + "\n这是一段你和用户的对话记忆，对其进行总结摘要，摘要使用第一人称'我'，并且提取其中的关键信息，以如下格式返回：\n 总结摘要 | 过去对话关键信息\n例如 用户张三问候我好，我礼貌回复，然后他问我langchain的向量库信息，我回答了他今年的问题，然后他又问了比特币价格。|Langchain, 向量库,比特币价格"),
                    ("user", "{input}")
                ]
            )

        
            chain = prompt | self.chatmodel
            summary = chain.invoke({"input": store_message, "who_you_are": Moods["default"]["roloSet"]})
            return summary
        except KeyError as e:
            print("总结出错")
            print(e)


    def get_memory(self, session_id: str = "session1"):
        """获取记忆历史，始终返回RedisChatMessageHistory对象"""
        try:
            # 🔧 修复：减少重复日志，使用logger而不是print
            logger.debug(f"获取记忆: session_id={session_id}")

            chat_message_history = RedisChatMessageHistory(
                session_id=session_id, url=redis_url)

            # 对超长的聊天记录进行摘要
            store_message = chat_message_history.messages
            if len(store_message) > 80:
                str_message = ""
                for message in store_message:
                    str_message += f"{type(message).__name__}: {message.content}\n"
                summary = self.summary_chain(str_message)
                chat_message_history.clear() #清空原有的对话
                chat_message_history.add_message(summary) #保存总结
                logger.info(f"对话摘要完成，session_id: {session_id}")
            else:
                logger.debug(f"聊天记录长度({len(store_message)})小于80，不进行摘要")

            # 始终返回RedisChatMessageHistory对象
            return chat_message_history
        except Exception as e:
            print(f"Error in get_memory: {e}")
            # 返回空的RedisChatMessageHistory对象而不是None
            return RedisChatMessageHistory(session_id=session_id, url=redis_url)

    def set_memory(self, session_id: str = "session1"):
        chat_memory = self.get_memory(session_id)
        if chat_memory is None:
            print("chat_memory is None")
            chat_memory = RedisChatMessageHistory(
                session_id=session_id, url=redis_url)
        
        # 下面是ConversationBufferMemory的参数解释：
        # llm: 指定用于对话的语言模型，这里传入self.chatmodel。
        # human_prefix: 人类用户在对话中的前缀标识，这里是"user"。
        # ai_prefix: AI助手在对话中的前缀标识，这里是"派大猩猩爱喝可乐"。
        # memory_key: 存储记忆内容的键名，这里用self.memoryKey。
        # output_key: 输出内容的键名，这里设为"output"。
        # return_messages: 是否返回消息对象而不是字符串，这里为True，返回消息对象列表。
        # max_token_limit: 对话记忆的最大token数，这里限制为3000。
        # chat_memory: 实际存储和管理对话历史的对象，这里传入chat_memory。
        self.memory = ConversationBufferMemory(
            llm = self.chatmodel,
            human_prefix = "user",
            ai_prefix= "派大猩猩爱喝可乐",
            memory_key=self.memoryKey,
            output_key="output",
            return_messages=True,
            max_token_limit = 3000,
            chat_memory= chat_memory
        )

        return self.memory

    def save_memory(self, session_id: str, memory_history):
        """保存记忆到Redis"""
        try:
            # memory_history 已经是 RedisChatMessageHistory 对象
            # 它会自动保存到Redis，无需额外操作
            logger.info(f"Memory saved for session: {session_id}")
            return True
        except Exception as e:
            logger.error(f"Failed to save memory: {e}")
            return False

    def clear_memory(self, session_id: str):
        """清除指定会话的记忆"""
        try:
            chat_memory = RedisChatMessageHistory(session_id=session_id, url=redis_url)
            chat_memory.clear()
            logger.info(f"Memory cleared for session: {session_id}")
            return True
        except Exception as e:
            logger.error(f"Failed to clear memory: {e}")
            return False

